#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <stdint.h>

#include <htslib/bgzf.h>
#include <htslib/sam.h>
#define read_len 150

int usage()
{
    fprintf(stderr,"bamSample [option] <in.bam>  <prefix> <bed>\n");
    fprintf(stderr,"-d  input bam depth\n");
	fprintf(stderr,"-D  output bam depth\n");
    return -1;
}

struct dot{
	int flag;
	int tid;
	int pos;
	char baminfo[10000];
	
};



int get_rand(int *a, int m, int n, int k)
{
    int i,j,t;
    for(i=0; i<k;)
    {
        t=rand()%(n-m+1)+m;
        for(j=0;j<i;j++)
            if(a[j]==t) break;
        if(j==i)
            a[i++]=t;
    }
        return 0;
}

int get_sign_rand(int *a, int m, int n, int k)
{
    int i,j,t;
    int tmp[1000];
    for(i=0;i<=m-n;i++)
    {
	tmp[i]=n+i;
    }
    for(i=0; i<m-n+1-k;)
    {
        t=rand()%(n-m+1)+m;
	if(tmp[t]==0 )
		break;
	else
	{
		tmp[t]=0;
		i++;
	}
    }
    j=0;
    for(i=0;i<=m-n;i++)
    {
	if(tmp[i]!=0)
	{
		a[j]=tmp[i];
		j++;
	}
	if(j>=k)	break;
    }

        return 0;
}



int get_next_rand(int *a, int m, int n, int k)
{
        int i;
        for(i=0;i<k;i++)
        {
                if(a[i]<n)
                        a[i]=a[i]+1;
                else
                        a[i]=0;
        }
        return 0;
}




int main( int argc , char *argv[])
{

	int train_sample_n=40;
	int test_sample_n=10;
	int check_sample_n=10;
	int i,c;
//	samFile *fp = NULL ;
	BGZF *fp =NULL;

	bam_hdr_t  *h = NULL ;
	bam1_t *b = NULL ;

   	int total_reads_num=0;
	
	int out_depth=50;
	int in_depth=1150;
	int expect=9;
	int expect_test=1;
	int expect_check=1;

    while( (c = getopt(argc,argv,"d:D:")) >=0 ){
	switch(c){
	    case 'd':  in_depth = atoi(optarg);  break ;
	    case 'D':  out_depth = atoi(optarg); break ;
	    default:
		printf("unKnow  argument\n");
		return -1;
	}
    }
    if( optind + 3 !=  argc)
	return usage();
	
	int sample_600x=in_depth*0.6*0.1;
	int sample_400x=in_depth*0.4*0.1;
/*
	expect      = out_depth*sample_600x/(in_depth*0.6);
	expect_check= out_depth*sample_600x/(in_depth*0.6);
	expect_test = out_depth*sample_400x/(in_depth*0.4);
*/

	expect      = out_depth*0.1;
        expect_check= out_depth*0.1;
        expect_test = out_depth*0.1;

		
	int min=0;
	int max_test  = sample_400x-1;
	int max_train = sample_600x-1;


	int rand_sign_en=0;
	if (expect > max_test*0.6 )
		rand_sign_en=1;

//    fp = sam_open(argv[optind],"r");
 	fp = bgzf_open( argv[optind], "r");
//   h = sam_hdr_read(fp);
	h = bam_hdr_read(fp);

//    hts_set_threads(fp,n_threads);
    b = bam_init1();

	BGZF *fp_train[50];
	BGZF *fp_test[10];
	BGZF *fp_check[10];
	char project[100];
        strcpy(project,argv[optind+1] );

	char train_Tbam[110][100];
        char test_Tbam[10][100];
	char check_Tbam[10][100];

	int train_rand[2000];
	int test_rand[2000];
	int check_rand[2000];

	char shell[1000];

	for(i=0;i<train_sample_n;i++)
	{
        	fp_train[i] = NULL ;
		sprintf(train_Tbam[i],"%s.%d.train.%d.bam", project,out_depth,i);
		fp_train[i] = bgzf_open( train_Tbam[i],"w");
		bam_hdr_write(fp_train[i], h);
	}
	for(i=0;i<10;i++)
        {
                fp_test[i] = NULL ;
		sprintf(test_Tbam[i],"%s.%d.test.%d.bam", project,out_depth,i);
                fp_test[i] = bgzf_open( test_Tbam[i],"w");
		bam_hdr_write(fp_test[i], h);
        }

	for(i=0;i<10;i++)
        {
                fp_check[i] = NULL ;
                sprintf(check_Tbam[i],"%s.%d.check.%d.bam", project,out_depth,i);
                fp_check[i] = bgzf_open( check_Tbam[i],"w");
                bam_hdr_write(fp_check[i], h);
        }

	
	FILE *fp_bed;
	fp_bed=fopen(argv[optind+2], "r");
	char bed_chr[10];
	int bed_chr_tid;
	int bed_pos;
	int bed_depth;
	int bam_read_en=bam_read1(fp,b);

	fscanf(fp_bed,"%s\t%d\t%*[^\n]\n", bed_chr, &bed_pos  );
	if(strcmp(bed_chr,"X")==0)
                bed_chr_tid=23;
        else if(strcmp(bed_chr,"Y")==0)
                bed_chr_tid=24;
        else if(strcmp(bed_chr,"M")==0)
                bed_chr_tid=22;
        else    bed_chr_tid=atoi(bed_chr)-1;

	int 	train_w_en=0;
	int 	test_w_en=0;


	int bed_read_cnt=0;


   while(bam_read_en >= 0 &&  !feof(fp_bed))
   {	

//	printf("%d\t%d\t%d\t%d\n", bed_chr_tid,bed_pos ,b->core.tid, b->core.pos);
	if (b->core.tid==bed_chr_tid && b->core.pos>=bed_pos-read_len && b->core.pos<=bed_pos+read_len)
	{
		bed_read_cnt++;
		total_reads_num++;
		if(total_reads_num%5==2||total_reads_num%5==4)
		{
				get_rand( test_rand, min, max_test, expect_test);

			for(i=0; i<expect_test; i++)
                        {
				if(test_rand[i]<10)
                                	test_w_en=bam_write1(fp_test[test_rand[i]], b);
                        }
			
		}
		else
		{
			get_rand( train_rand, min, max_train, expect);
			for(i=0; i<expect; i++)
			{	
				if(train_rand[i]<train_sample_n)
					train_w_en=bam_write1(fp_train[train_rand[i]], b);
			}

			get_rand( check_rand, min, max_train, expect_check);		
			for(i=0; i<expect_check; i++)
                        {
				if(check_rand[i]<10)
	                                train_w_en=bam_write1(fp_check[check_rand[i]], b);
                        }
		}
		bam_read_en=bam_read1(fp,b);
	}
	else if(b->core.tid<bed_chr_tid || (b->core.tid<=bed_chr_tid && b->core.pos<bed_pos-read_len))
	{
		bam_read_en=bam_read1(fp,b);
	}
	else
	{
	//	printf("%s\t%d\t%d\t%d\n",bed_chr, bed_pos, bed_depth,bed_read_cnt);
		bed_read_cnt=0;
		fscanf(fp_bed,"%s\t%d\t%*[^\n]\n", bed_chr, &bed_pos  );
		if(strcmp(bed_chr,"X")==0)
	                bed_chr_tid=23;
        	else if(strcmp(bed_chr,"Y")==0)
                	bed_chr_tid=24;
	        else if(strcmp(bed_chr,"M")==0)
        	        bed_chr_tid=22;
	        else    bed_chr_tid=atoi(bed_chr)-1;
	}

    }

//    printf("%d\n", total_reads_num );   
	for(i=0;i<train_sample_n;i++)
		bgzf_close(fp_train[i]);
	for(i=0;i<10;i++)
		bgzf_close(fp_test[i]);

	for(i=0;i<10;i++)
                bgzf_close(fp_check[i]);
    bgzf_close(fp);
    bam_hdr_destroy(h);
    bam_destroy1(b);
    return 0 ;
}
